import os

from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms

from config import parser_config
from utils import WaterMark

data_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])


class Dataset(Dataset):
    def __init__(self, type="train"):
        super(Dataset, self).__init__()
        self.args = parser_config()
        self.imageList(type)
        self.waterMarker = WaterMark(0.35)
        # self.pepperNoise = AddPepperNoise(0.999)

    def __len__(self):
        return self.length

    def __getitem__(self, x):
        data = self.images[x]
        data = Image.open(data).convert('RGB')
        data = data.resize((256, 256))
        target = data
        data = self.waterMarker(data)
        # data = self.pepperNoise(data)

        data = data_transform(data)
        target = data_transform(target)
        return data, target

    def imageList(self, type):
        if type == "train":
            self.images = test_images(self.args.train_data_path)

        else:
            self.images = test_images(self.args.test_data_path)
        self.length = len(self.images)


def train_images(path):
    imagePath = []
    for file in os.listdir(path):
        imageFile = os.path.join(path, file).replace('\\','/')
        for image in os.listdir(imageFile):
            imagePath.append(os.path.join(imageFile, image).replace('\\','/'))

    return imagePath


def test_images(path):
    imagePath = []
    for file in os.listdir(path):
        imagePath.append(os.path.join(path, file).replace('\\','/'))
    return imagePath

